import pymesh
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

# 创建一个四面体圆柱网格
vertices = np.array([
    [0, 0, 0],
    [1, 0, 0],
    [0, 1, 0],
    [0, 0, 1],
    [0, 0, 2],
    [1, 0, 2],
    [0, 1, 2],
    [0, 0, 3]
])
faces = np.array([
    [0, 1, 2],
    [0, 1, 3],
    [0, 2, 3],
    [1, 2, 3],
    [4, 5, 6],
    [4, 5, 7],
    [4, 6, 7],
    [5, 6, 7],
    [0, 1, 4],
    [1, 4, 5],
    [1, 2, 5],
    [2, 5, 6],
    [2, 3, 6],
    [3, 6, 7],
    [3, 0, 7],
    [0, 4, 7]
])

mesh = pymesh.form_mesh(vertices, faces)

# 可视化四面体圆柱网格
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
ax.plot_trisurf(mesh.vertices[:, 0], mesh.vertices[:, 1], mesh.vertices[:, 2], triangles=mesh.faces)
ax.set_xlabel('X')
ax.set_ylabel('Y')
ax.set_zlabel('Z')
plt.show()

# 初始化基础属性
mesh.add_attribute("alpha_l")  # 液相含率
mesh.add_attribute("alpha_g")  # 气相含率
mesh.add_attribute("velocity_l")  # 液相速度
mesh.add_attribute("velocity_g")  # 气相速度
mesh.add_attribute("pressure")  # 压力
mesh.add_attribute("enthalpy_l")  # 液相焓
mesh.add_attribute("enthalpy_g")  # 气相焓
mesh.add_attribute("mass")  # 质量
mesh.add_attribute("density")  # 混合物密度
mesh.add_attribute("flow")  # 总流量
mesh.add_attribute("momentum")  # 动量
mesh.add_attribute("mass_exchange")  # 质量交换项
mesh.add_attribute("flow_l")  # 液相流量
mesh.add_attribute("flow_v")  # 气相流量
mesh.add_attribute("density_l")  # 液相密度
mesh.add_attribute("density_v")  # 气相密度
mesh.add_attribute("gravity_x")  # 重力加速度
mesh.add_attribute("drag_coefficient")  # 阻力系数
mesh.add_attribute("heat_flux_wall_liquid")  # 墙面对液体的热流
mesh.add_attribute("heat_flux_wall_vapor")  # 墙面对气体的热流
mesh.add_attribute("heat_flux_liquid")  # 液相之间的热流

# 初始化时间和空间步长
time_step = 0.01  # 时间步长 (单位: s)
x_step = 0.1      # 空间步长 (单位: m)
A = 1.0           # 面积 (单位: m²)

# 初始条件
mesh.set_attribute("alpha_l", np.full(mesh.num_faces, 0.5))  # 液相含率
mesh.set_attribute("alpha_g", np.full(mesh.num_faces, 0.5))  # 气相含率
mesh.set_attribute("velocity_l", np.full(mesh.num_faces, 1.0))  # 液相速度
mesh.set_attribute("velocity_g", np.full(mesh.num_faces, 0.1))  # 气相速度
mesh.set_attribute("pressure", np.full(mesh.num_faces, 1e5))  # 压力 (单位: Pa)
mesh.set_attribute("enthalpy_l", np.full(mesh.num_faces, 1e4))  # 液相焓 (单位: J/kg)
mesh.set_attribute("enthalpy_g", np.full(mesh.num_faces, 2e4))  # 气相焓 (单位: J/kg)
mesh.set_attribute("mass", np.full(mesh.num_faces, 1.0))  # 质量
mesh.set_attribute("density", np.full(mesh.num_faces, 1.2))  # 混合物密度 (单位: kg/m³)
mesh.set_attribute("flow", np.full(mesh.num_faces, 0.1))  # 总流量 (单位: kg/s)
mesh.set_attribute("momentum", np.zeros(mesh.num_faces))  # 动量 (初始化为0)
mesh.set_attribute("mass_exchange", np.full(mesh.num_faces, 0.0))  # 质量交换项
mesh.set_attribute("flow_l", np.full(mesh.num_faces, 0.05))  # 液相流量 (单位: kg/s)
mesh.set_attribute("flow_v", np.full(mesh.num_faces, 0.05))  # 气相流量 (单位: kg/s)
mesh.set_attribute("density_l", np.full(mesh.num_faces, 1000.0))  # 液相密度 (单位: kg/m³)
mesh.set_attribute("density_v", np.full(mesh.num_faces, 0.6))  # 气相密度 (单位: kg/m³)
mesh.set_attribute("gravity_x", np.full(mesh.num_faces, 9.81))  # 轴向重力加速度 (单位: m/s²)
mesh.set_attribute("drag_coefficient", np.full(mesh.num_faces, 0.1))  # 阻力系数
mesh.set_attribute("heat_flux_wall_liquid", np.full(mesh.num_faces, 0.0))  # 墙面对液体的热流
mesh.set_attribute("heat_flux_wall_vapor", np.full(mesh.num_faces, 0.0))  # 墙面对气体的热流
mesh.set_attribute("heat_flux_liquid", np.full(mesh.num_faces, 0.0))  # 液相之间的热流

# 涉及到的中间变量
mesh.add_attribute("W_l_prime")  # 液相质量交换速率的导数项
mesh.add_attribute("W_v_prime")  # 气相质量交换速率的导数项
mesh.add_attribute("delta_U_l")  # 液相速度变化量
mesh.add_attribute("delta_U_v")  # 气相速度变化量
mesh.add_attribute("Psi_l_n")    # 热流相关系数
mesh.add_attribute("delta_h_l")  # 液相焓变化量
mesh.add_attribute("delta_h_v")  # 气相焓变化量

# 初始化中间变量
mesh.set_attribute("W_l_prime", np.full(mesh.num_faces, 0.0))
mesh.set_attribute("W_v_prime", np.full(mesh.num_faces, 0.0))
mesh.set_attribute("delta_U_l", np.full(mesh.num_faces, 0.0))
mesh.set_attribute("delta_U_v", np.full(mesh.num_faces, 0.0))
mesh.set_attribute("Psi_l_n", np.full(mesh.num_faces, 0.0))
mesh.set_attribute("delta_h_l", np.full(mesh.num_faces, 0.0))
mesh.set_attribute("delta_h_v", np.full(mesh.num_faces, 0.0))
def mass_conservation(mesh, dt, dx):
    # 质量守恒方程的离散化
    rho = mesh.get_attribute("density")
    F = mesh.get_attribute("flow")
    W = mesh.get_attribute("mass_exchange")
    
    # 计算质量变化率
    drho_dt = (rho[1:] - rho[:-1]) / dt
    dF_dx = (F[1:] - F[:-1]) / dx

    mass_balance = drho_dt + dF_dx + W[:-1]
    return mass_balance


def axial_momentum_conservation(mesh, dt, dx, A, K_U):
    # 轴向动量守恒方程的离散化
    F = mesh.get_attribute("flow")
    P = mesh.get_attribute("pressure")
    rho = mesh.get_attribute("density")
    g_x = mesh.get_attribute("gravity_x")
    
    F_l = mesh.get_attribute("flow_l")
    F_v = mesh.get_attribute("flow_v")
    U_l = mesh.get_attribute("velocity_l")
    U_v = mesh.get_attribute("velocity_v")
    
    W_l = mesh.get_attribute("mass_exchange_l")
    W_l_prime = mesh.get_attribute("W_l_prime")
    delta_U_l = mesh.get_attribute("delta_U_l")
    W_v_prime = mesh.get_attribute("W_v_prime")
    delta_U_v = mesh.get_attribute("delta_U_v")
    
    dF_dt = (F[1:] - F[:-1]) / dt
    dF_dx = (F_l[1:] * U_l[1:] + F_v[1:] * U_v[1:] - F_l[:-1] * U_l[:-1] - F_v[:-1] * U_v[:-1]) / dx
    dP_dx = (P[1:] - P[:-1]) / dx
    
    momentum_balance = (
        dF_dt + dF_dx
        + dP_dx * A
        - W_l[:-1] * U_l[:-1]
        - W_v_prime[:-1] * delta_U_v[:-1]
        - K_U * np.abs(F[:-1]) * F[:-1]
        + rho[:-1] * g_x[:-1]
    )
    return momentum_balance


def transverse_momentum_conservation(mesh, dt, dx, S, K_v):
    # 横向动量守恒方程的离散化
    rho = mesh.get_attribute("density")
    V = mesh.get_attribute("velocity_v")
    F = mesh.get_attribute("flow")
    P = mesh.get_attribute("pressure")
    W = mesh.get_attribute("mass_exchange")
    W_v_prime = mesh.get_attribute("W_v_prime")
    delta_U_v = mesh.get_attribute("delta_U_v")
    g_y = 9.81  # 重力在y方向的加速度
    
    dV_dt = (V[1:] - V[:-1]) / dt
    dFV_dx = (F[1:] * V[1:] - F[:-1] * V[:-1]) / dx
    dF_dx = (F[1:] - F[:-1]) / dx
    dP_dy = (P[1:] - P[:-1]) / dx  # 近似在x方向
    
    transverse_balance = (
        9 * rho[:-1] * dV_dt + dFV_dx - V[:-1] * dF_dx
        - 9 * dP_dy
        + S * dx * K_v * W[:-1] * W[:-1]
        + (W[:-1] * delta_U_v[:-1] + W_v_prime[:-1] * delta_U_v[:-1]) / (S * dx)
        + 9 * rho[:-1] * g_y
    )
    return transverse_balance


def enthalpy_conservation_liquid(mesh, dt, dx, A):
    # 液相焓守恒方程的离散化
    alpha = mesh.get_attribute("alpha_l")
    rho_l = mesh.get_attribute("density_l")
    h_l = mesh.get_attribute("enthalpy_l")
    F_l = mesh.get_attribute("flow_l")
    W_l = mesh.get_attribute("mass_exchange_l")
    
    dh_l_dt = (h_l[1:] - h_l[:-1]) / dt
    dFh_dx = (F_l[1:] * h_l[1:] - F_l[:-1] * h_l[:-1]) / dx
    dF_dx = (F_l[1:] - F_l[:-1]) / dx
    
    enthalpy_balance_l = (
        (1 - alpha[:-1]) * A * dx * rho_l[:-1] * dh_l_dt
        + dx * dFh_dx - dx * h_l[:-1] * dF_dx + W_l[:-1] * h_l[:-1]
    )
    return enthalpy_balance_l


def enthalpy_conservation_vapor(mesh, dt, dx, A):
    # 汽相焓守恒方程的离散化
    alpha = mesh.get_attribute("alpha_v")
    rho_v = mesh.get_attribute("density_v")
    h_v = mesh.get_attribute("enthalpy_v")
    F_v = mesh.get_attribute("flow_v")
    W_v = mesh.get_attribute("mass_exchange_v")
    
    dh_v_dt = (h_v[1:] - h_v[:-1]) / dt
    dFh_dx = (F_v[1:] * h_v[1:] - F_v[:-1] * h_v[:-1]) / dx
    dF_dx = (F_v[1:] - F_v[:-1]) / dx
    
    enthalpy_balance_v = (
        alpha[:-1] * A * dx * rho_v[:-1] * dh_v_dt
        + dx * dFh_dx - dx * h_v[:-1] * dF_dx + W_v[:-1] * h_v[:-1]
    )
    return enthalpy_balance_v


def mixture_enthalpy_conservation(mesh, dt, dx, A):
    # 混合物焓守恒方程的离散化
    rho = mesh.get_attribute("density")
    h = mesh.get_attribute("enthalpy")
    F_l = mesh.get_attribute("flow_l")
    h_l = mesh.get_attribute("enthalpy_l")
    F_v = mesh.get_attribute("flow_v")
    h_v = mesh.get_attribute("enthalpy_v")
    W_l = mesh.get_attribute("mass_exchange_l")
    W_v = mesh.get_attribute("mass_exchange_v")
    
    dh_dt = (h[1:] - h[:-1]) / dt
    dFh_dx = (F_l[1:] * h_l[1:] + F_v[1:] * h_v[1:] - F_l[:-1] * h_l[:-1] - F_v[:-1] * h_v[:-1]) / dx
    dF_dx = (F_l[1:] + F_v[1:] - F_l[:-1] - F_v[:-1]) / dx
    
    enthalpy_balance_mixture = (
        A * dx * rho[:-1] * dh_dt
        + dx * dFh_dx - dx * h[:-1] * dF_dx
        + (W_l[:-1] * h_l[:-1] + W_v[:-1] * h_v[:-1])
    )
    return enthalpy_balance_mixture


